`include "defines.svh"

module StageWB (
    input   logic                           clock,
    input   logic                           reset,

    // for debug pipeline
    output  logic [63:0]                    debug_wb_pc,
    output  logic                           debug_wb_valid,

    // StageWB -> Regfile
    output  logic [`WB2RF_BUS_WD-1:0]       wb2rf_bus,

    // StageWB -> CSR
    output  logic [`WB2CSR_BUS_WD-1:0]      wb2csr_bus,

    // CSR -> StageWB
    input   logic                           csr_taken,

    // StageWB <- StageMA
    input   logic                           ma2wb_valid,
    input   logic [`MA2WB_BUS_WD-1:0]       ma2wb_bus,

    // StageWB -> StageMA
    output  logic                           wb2ma_allow,

    // StageWB -> StageIF/EX @ Cancel Flag
    output  logic                           wb_cancel,

    // Forward & Block from StageWB to StageID
    output  logic [`WB_RF_HZD_BUS_WD-1:0]   wb_rf_hzd_bus,
    output  logic [`WB_CSR_BLK_BUS_WD-1:0]  wb_csr_blk_bus,

    // Commit
    output  logic [63:0]                    icache_miss_count,
    output  logic [63:0]                    dcache_miss_count,
    output  logic [63:0]                    commit_inst_count,
    output  logic [63:0]                    mmio_inst_count,
    output  logic [63:0]                    mem_inst_count,
    output  logic                           commit_valid,
    output  logic [63:0]                    commit_pc,
    output  logic [31:0]                    commit_inst,
    output  logic                           commit_skip_ref
);

    logic           wb_ready;
    logic           wb_valid;

    assign debug_wb_pc = wb_pc;
    assign debug_wb_valid = wb_valid;

/* ************************* Pipeline ************************* */
/* -------------------- StageMA <=> StageWB -------------------- */
    // for debug signal
    logic           inst_mem;
    logic           skip_ref;
    logic           icache_miss;
    logic           dcache_miss;
    // to StageWB signal
    logic [63:0]    wb_pc;
    logic [31:0]    wb_inst;
    logic [63:0]    wb_next_pc;
    logic [15:0]    ma2wb_excp_flags;
    logic [15:0]    ma2wb_intr_flags;
    // to Regfile signal
    logic           ma2wb_rf_we;
    logic [4:0]     wb_rf_waddr;
    logic [63:0]    wb_rf_wdata;
    // to CSR signal
    logic           inst_mret;
    logic           ma2wb_csr_we;
    logic [11:0]    wb_csr_waddr;
    logic [63:0]    wb_csr_wdata;

    logic [`MA2WB_BUS_WD-1:0]   ma2wb_bus_R;

    assign {
        // for debug
        inst_mem,
        skip_ref,
        icache_miss,
        dcache_miss,
        // to StageWB
        wb_pc,
        wb_inst,
        wb_next_pc,
        ma2wb_excp_flags,
        ma2wb_intr_flags,
        // to Regfile
        ma2wb_rf_we,
        wb_rf_waddr,
        wb_rf_wdata,
        // to CSR
        inst_mret,
        ma2wb_csr_we,
        wb_csr_waddr,
        wb_csr_wdata
    } = ma2wb_bus_R;

    assign  wb_ready = 1'b1;
    assign  wb2ma_allow = (~wb_valid) | (wb_ready);

    always_ff @ (posedge clock) begin
        if (reset | csr_taken) begin
            wb_valid <= 1'b0;
        end
        else if (wb2ma_allow) begin
            wb_valid <= ma2wb_valid;
        end
    end

    always_ff @ (posedge clock) begin
        if (reset) begin
            ma2wb_bus_R <= `MA2WB_BUS_WD'b0;
        end else
        if (ma2wb_valid & wb2ma_allow) begin
            ma2wb_bus_R <= ma2wb_bus;
        end
    end
/* -------------------- StageMA <=> StageWB -------------------- */
/* ************************* Pipeline ************************* */


/* ******************************** RegFile ******************************** */
    wire wb_rf_we = ma2wb_rf_we & wb_valid;
    assign wb2rf_bus = { wb_rf_we, wb_rf_waddr, wb_rf_wdata };

    // RegFile Forward & Block
    logic           wb_rf_blk_flag, wb_rf_fwd_flag;
    logic [63:0]    wb_rf_fwd_data;
    assign wb_rf_blk_flag = 1'b0;
    assign wb_rf_fwd_flag = wb_valid & ma2wb_rf_we & (|wb_rf_waddr);
    assign wb_rf_fwd_data = wb_rf_wdata;
    assign wb_rf_hzd_bus = {
        wb_rf_blk_flag,
        wb_rf_fwd_flag,
        wb_rf_waddr,
        wb_rf_fwd_data
    };
/* ******************************** RegFile ******************************** */


/* ************************* CSR ************************* */
    // CSR Block
    logic wb_csr_blk_flag;
    assign wb_csr_blk_flag = wb_valid & wb_csr_we;
    assign wb_csr_blk_bus = {
        wb_csr_blk_flag,
        wb_csr_waddr
    };

    // to CSR
    wire wb_csr_we = ma2wb_csr_we & wb_valid;
    wire wb_mret_taken = wb_has_mret;
    wire wb_excp_taken = wb_has_excp;
    wire wb_intr_taken = wb_has_intr;
    wire wb_trap_taken = wb_excp_taken | wb_intr_taken;

    logic [63:0]    wb_trap_cause;
    Encoder #(16, 4)    cause_code_encoder (
        .in(wb_excp_taken ? ma2wb_excp_flags : ma2wb_intr_flags),
        .out(wb_trap_cause[3:0])
    );
    assign wb_trap_cause[62:4] = 59'b0;
    assign wb_trap_cause[63] = wb_intr_taken;

    wire [63:0] wb_trap_pc = wb_excp_taken ? wb_pc : wb_next_pc;
    wire [31:0] wb_trap_inst = wb_inst;
    assign wb2csr_bus = {
        wb_csr_we,
        wb_csr_waddr,
        wb_csr_wdata,
        wb_mret_taken,
        wb_trap_taken,
        wb_trap_cause,
        wb_trap_pc,
        wb_trap_inst
    };
/* ************************* CSR ************************* */


/* ************************************* Trap ************************************* */
    wire wb_has_mret = wb_valid & inst_mret;
    wire wb_has_excp = wb_valid & (|ma2wb_excp_flags);
    wire wb_has_intr = wb_valid & (|ma2wb_intr_flags);
/* ----------------------------- Cancel Flag ----------------------------- */
    assign wb_cancel = wb_has_excp | wb_has_intr | wb_has_mret;
/* ----------------------------- Cancel Flag ----------------------------- */
/* ************************************* Trap ************************************* */


/* ********************* Commit ********************* */
    always_ff @(posedge clock) begin
        if (reset) begin
            commit_valid    <= 0;
            commit_pc       <= 0;
            commit_inst     <= 0;
            commit_skip_ref <= 0;
        end
        else begin
            commit_valid    <= wb_valid;
            commit_pc       <= wb_pc;
            commit_inst     <= wb_inst;
            commit_skip_ref <= skip_ref;
        end
    end
/* ********************* Commit ********************* */


/* ********************* Perf Counter ********************* */
    always_ff @(posedge clock) begin
        if (reset) begin
            icache_miss_count <= 0;
            dcache_miss_count <= 0;
            commit_inst_count <= 0;
            mmio_inst_count <= 0;
            mem_inst_count <= 0;
        end else begin
            if (wb_valid) begin
                commit_inst_count <= commit_inst_count + 1;
            end
            if (wb_valid & icache_miss) begin
                icache_miss_count <= icache_miss_count + 1;
            end
            if (wb_valid & dcache_miss) begin
                dcache_miss_count <= dcache_miss_count + 1;
            end
            if (wb_valid & skip_ref) begin
                mmio_inst_count <= mmio_inst_count + 1;
            end
            if (wb_valid & inst_mem) begin
                mem_inst_count <= mem_inst_count + 1;
            end
        end
    end
/* ********************* Perf Counter ********************* */

endmodule
